import os
import time
import logging
import threading
from datetime import datetime
from flask import jsonify, request, app
from kubernetes import client, config
from utils.dingtalk import send_ding_ding
from utils.env_processing import token_required


def test_k8s_connection():
    """
    Kubernetes权限测试用例
    :return: True if successful, False otherwise
    """
    try:
        config.load_incluster_config()
    except Exception as e:
        logging.error(f"Failed to load in-cluster configuration: {e}")
        return False

    v1 = client.CoreV1Api()
    logging.info("This is a test case that lists the kube-system pods:")
    try:
        pods = v1.list_namespaced_pod('kube-system')
        for pod in pods.items:
            logging.info(pod.metadata.name)
        return True
    except Exception as e:
        logging.error(f"Failed to list pods in kube-system namespace: {e}")
        return False


def get_api_client(resource_type: str):
    """
    获取Kubernetes API客户端
    :param resource_type: 资源类型
    :return: API客户端实例
    """
    # 加载集群配置
    config.load_incluster_config()

    api_mapping = {
        'deployment': client.AppsV1Api,
        'stateful_set': client.AppsV1Api,
        'daemon_set': client.AppsV1Api,
        'cron_job': client.BatchV1Api,
        'job': client.BatchV1Api,
    }

    if resource_type not in api_mapping:
        raise ValueError('Invalid resource type')

    return api_mapping[resource_type]()


def get_resource_status(namespace: str, resource_type: str, resource_name: str):
    """
    获取当前资源状态，将状态信息推送至钉钉。
    :param namespace: 资源所在的命名空间
    :param resource_type: 资源类型
    :param resource_name: 资源名称
    :return:
    """
    try:
        api = get_api_client(resource_type)
        core_v1 = client.CoreV1Api()

        while True:
            try:
                # 获取资源对象
                resource_method = getattr(api, f'read_namespaced_{resource_type}')
                current = resource_method(resource_name, namespace)

                # 获取 Pod 列表
                if resource_type == 'deployment':
                    label_selector = ",".join([f"{k}={v}" for k, v in current.spec.selector.match_labels.items()])
                    pods = core_v1.list_namespaced_pod(namespace=namespace, label_selector=label_selector)
                else:
                    # 对于其他资源类型，直接获取其 Pod 列表
                    pods = core_v1.list_namespaced_pod(namespace=namespace, label_selector=f"app={resource_name}")

                # 检查每个 Pod 中的容器状态
                all_running = True
                for pod in pods.items:
                    logging.info(f"Checking Pod: {pod.metadata.name}")
                    for container_status in pod.status.container_statuses:
                        state = container_status.state
                        container_name = container_status.name
                        logging.info(f"Container: {container_name}")

                        if state.waiting is not None:
                            reason = state.waiting.reason
                            logging.info(f"State: Waiting, Reason: {reason}")
                            if reason not in ['ContainerCreating', 'PodInitializing']:
                                status = reason
                                logging.info(f"Sending message for Pod {pod.metadata.name}: {status}")
                                push_status = send_ding_ding(resource_name, status)
                                if push_status == 0:
                                    return
                                all_running = False
                                break
                        elif state.terminated is not None:
                            logging.info(f"State: Terminated, Exit Code: {state.terminated.exit_code}")
                            status = 'Terminated'
                            logging.info(f"Sending message for Pod {pod.metadata.name}: {status}")
                            push_status = send_ding_ding(resource_name, status)
                            if push_status == 0:
                                return
                            all_running = False
                            break
                        elif state.running is not None:
                            logging.info(f"State: Running")
                        else:
                            logging.info(f"State: Unknown")
                            status = 'Unknown'
                            logging.info(f"Sending message for Pod {pod.metadata.name}: {status}")
                            push_status = send_ding_ding(resource_name, status)
                            if push_status == 0:
                                return
                            all_running = False
                            break

                    if not all_running:
                        break  # 如果发现异常状态，退出当前 Pod 的循环

                # 如果所有容器都处于运行状态，发送成功消息
                if all_running:
                    status = 'Available'
                    logging.info(f"All containers are running, sending message: {status}")
                    push_status = send_ding_ding(resource_name, status)
                    if push_status == 0:
                        return
                    break  # 退出循环

            except client.ApiException as e:
                logging.error(f'Failed to obtain resources: {resource_name} {e.reason}')
                break  # 如果获取失败，退出循环

            # 等待30秒后再次检查
            time.sleep(30)

    except ValueError as e:
        logging.error(str(e))


@token_required
def update_k8s_image():
    """
    处理请求的JSON，用于更新资源镜像
    对于多容器资源，必须指定容器名称，否则返回错误
    :return:
    """
    # 解析请求数据
    data = request.get_json()
    resource_name = data['name']
    namespace = data['namespace']
    image = data['image']
    resource_type = data['resource_type']
    container_name = data.get('container', None)

    try:
        api = get_api_client(resource_type)
        resource_method = getattr(api, f'read_namespaced_{resource_type}')
        current = resource_method(resource_name, namespace)
    except (client.ApiException, ValueError) as e:
        return jsonify({'error': str(e)}), 400

    # 检查容器数量
    container_count = len(current.spec.template.spec.containers)
    if container_count > 1 and not container_name:
        return jsonify({
            'error': 'Container name must be specified for multi-container resources',
            'containers': [c.name for c in current.spec.template.spec.containers]
        }), 400

    # 查找并更新指定容器
    updated = False
    for container in current.spec.template.spec.containers:
        if container_name is None or container.name == container_name:
            container.image = image
            updated = True

    if not updated:
        return jsonify({
            'error': f'Container {container_name} not found',
            'available_containers': [c.name for c in current.spec.template.spec.containers]
        }), 400

    try:
        patch_method = getattr(api, f'patch_namespaced_{resource_type}')
        result = patch_method(resource_name, namespace, current)
    except client.ApiException as e:
        return jsonify({'error': e.reason, 'message': e.body}), 400

    logging.debug(result.to_dict() if result else None)

    # 获取push环境变量，默认为disable
    push_enabled = os.getenv('push', 'disable')
    if push_enabled.lower() == 'enable':
        time.sleep(5)
        threading.Thread(target=get_resource_status, args=(namespace, resource_type, resource_name)).start()

    return jsonify({'message': 'Image update successful'}), 200


def update_k8s_image_form():
    """
    处理请求的表单数据，用于更新资源镜像
    对于多容器资源，必须指定容器名称，否则返回错误
    :return:
    """
    # 解析请求数据
    data = request.form
    project = data.get('project')
    namespace = data.get('namespace')
    image = data.get('image')
    container_name = data.get('container', None)
    resource_type = data.get('resource_type')

    try:
        api = get_api_client(resource_type)
        resource_method = getattr(api, f'read_namespaced_{resource_type}')
        current = resource_method(name=project, namespace=namespace)
    except (client.ApiException, ValueError) as e:
        return jsonify({'error': str(e)}), 400

    # 检查容器数量
    container_count = len(current.spec.template.spec.containers)
    if container_count > 1 and not container_name:
        return jsonify({
            'error': 'Container name must be specified for multi-container resources',
            'containers': [c.name for c in current.spec.template.spec.containers]
        }), 400

    # 查找并更新指定容器
    updated = False
    for container in current.spec.template.spec.containers:
        if container_name is None or container.name == container_name:
            container.image = image
            updated = True

    if not updated:
        return jsonify({
            'error': f'Container {container_name} not found',
            'available_containers': [c.name for c in current.spec.template.spec.containers]
        }), 400

    try:
        patch_method = getattr(api, f'patch_namespaced_{resource_type}')
        result = patch_method(name=project, namespace=namespace, body=current)
    except client.ApiException as e:
        return jsonify({'error': e.reason, 'message': e.body}), 400

    logging.debug(result.to_dict() if result else None)

    return jsonify({'message': 'Image update successful'}), 200


@token_required
def rollback_deployment(namespace: str, deployment_name: str, revision: int = None):
    try:
        config.load_incluster_config()
        apps_v1 = client.AppsV1Api()

        # 检查部署是否存在
        try:
            deployment = apps_v1.read_namespaced_deployment(
                name=deployment_name,
                namespace=namespace
            )
        except client.ApiException as e:
            if e.status == 404:
                return False, f"Deployment {deployment_name} not found in namespace {namespace}"
            raise

        # 获取关联的ReplicaSet（用于版本回滚）
        label_selector = f"app={deployment_name}"
        replica_sets = apps_v1.list_namespaced_replica_set(
            namespace=namespace,
            label_selector=label_selector
        ).items

        if revision:
            # 查找指定版本的ReplicaSet
            target_rs = next(
                (rs for rs in replica_sets if
                 rs.metadata.annotations.get('deployment.kubernetes.io/revision') == str(revision)),
                None
            )

            if not target_rs:
                return False, f"Revision {revision} not found for deployment {deployment_name}"

            # 执行版本回滚
            body = {
                "spec": {
                    "template": target_rs.spec.template,
                    "revisionHistoryLimit": deployment.spec.revision_history_limit
                }
            }
            message = f"Rollback to revision {revision} succeeded"
        else:
            # 普通回滚（通过修改注解触发重新部署）
            body = {
                "spec": {
                    "template": {
                        "metadata": {
                            "annotations": {
                                "kubectl.kubernetes.io/restartedAt": datetime.now().isoformat()
                            }
                        }
                    }
                }
            }
            message = "Rollback initiated (new deployment triggered)"

        apps_v1.patch_namespaced_deployment(
            name=deployment_name,
            namespace=namespace,
            body=body
        )

        return True, message

    except Exception as e:
        logging.error(f"Rollback failed: {str(e)}")
        return False, f"Rollback failed: {str(e)}"
